#include "channeler.h"
#include "data.h"
#include <stddef.h>
#include <stdio.h>
#include "psrfold.h"
#include "variables.h"

void pfb_ring(double *data,fftw_complex *subband,double *coeffs,int index,int ntap,int nfft)
{
    // double *sum = calloc(sizeof(double),nfft);
    fftw_complex *sum_c = calloc(sizeof(fftw_complex) , nfft);
    double *temp = malloc(sizeof(double) * ntap);
    /* index belong to [0,ntap) */
    index %= ntap;
    for(int i = 0;i < nfft;++i){
        for(int j = index;j < ntap;++j){
            temp[j-index] = data[j*nfft+i];
            // sum[i] += data[j * nfft + i] * coeffs[(j-index) * nfft + i];
        }
        for(int j = 0;j < index;++j){
            temp[j+ntap-index] = data[j*nfft+i];
            // sum[i] += data[j * nfft + i] * coeffs[(ntap+j-index) * nfft + i];
        }
        for(int j = 0;j < ntap;++j){
            sum_c[i][0] += temp[ntap - j - 1] * coeffs[j* nfft + i];
        }
    }
    // real2complex(sum,sum_c,nfft);
    fft_c2c(sum_c,subband,nfft);
    fftw_free(sum_c);
    free(temp);
    // free(sum);
}

void real2complex(double *in, fftw_complex *out,long size)
{
    for(int i = 0;i < size;++i){
        out[i][0] = in[i];
        out[i][1] = 0;
    }
}

void insert_data(double *in,double *data,int nfft,int ntap,int index)
{
    int iindex = index % ntap;
    for(int i = 0;i < nfft;++i){
        data[iindex * nfft + i] = in[index * nfft + i];
    }
}

void fft_c2c(fftw_complex *in,fftw_complex *out,long size)
{
    fftw_plan_with_nthreads(4);
    fftw_plan plan = fftw_plan_dft_1d(size,in,out,FFTW_BACKWARD,FFTW_ESTIMATE);
    if(plan == NULL) {
        printf("ifft:Create fftw_plan error!");
        exit(1);
    }
    fftw_execute(plan);
    fftw_destroy_plan(plan);
    fftw_cleanup_threads();
}

void fft_r2r(double *in,double *out,int size)
{
    fftw_plan_with_nthreads(4);
    fftw_r2r_kind kind = FFTW_REDFT00; // DCT-1
    fftw_plan plan = fftw_plan_r2r_1d(size,in,out,kind,FFTW_ESTIMATE);
    if(plan == NULL) {
        printf("fft:Create fftw_plan error!");
        exit(1);
    }
    fftw_execute(plan);
    fftw_destroy_plan(plan);
    fftw_cleanup_threads();
}

void fft_c2r(fftw_complex *in,double *out,long size)
{
    fftw_plan_with_nthreads(4);
    fftw_plan plan = fftw_plan_dft_c2r_1d(size,in,out,FFTW_ESTIMATE);
    if(plan == NULL) {
        printf("ifft:Create fftw_plan error!");
        exit(1);
    }
    fftw_execute(plan);
    fftw_destroy_plan(plan);
    fftw_cleanup_threads();
}

void fft_r2c(double *in,fftw_complex *out,int size)
{
    fftw_plan_with_nthreads(4);
    fftw_plan plan = fftw_plan_dft_r2c_1d(size,in,out,FFTW_ESTIMATE);
    if(plan == NULL) {
        printf("fft:Create fftw_plan error!");
        exit(1);
    }
    fftw_execute(plan);
    fftw_destroy_plan(plan);
    fftw_cleanup_threads();
}

// void channeler(double *data,double *coeffs,fftw_complex *subband,
//                 double *buffer,int *res_num,int ntap,int nfft)
// {
//     if(*res_num == 0){ return;}
//     int num = *res_num;
//     // printf("num = %d\n",num);
//     for(int i = 0;i < num;++i){
//         insert_data(data,buffer,nfft,ntap,i);
//         pfb_ring(buffer,subband+i*nfft,coeffs,i+1,ntap,nfft);
//     }
// }

void extract_data(double *data,double *out,size_t len,int index,int step1,int step2)
{
    for(size_t i = index,j = 0;i < len;){
        out[j] = data[i];
        j = j + step1;
        i = i + step2;
    }
}

void channeler(double *data,double *coeffs,fftw_complex *subband,
                long res_num,int ntap,int nfft)
{
    double *sdata = malloc(sizeof(double) * res_num);
    double *scoeff = malloc(sizeof(double) * ntap);

    double *out = malloc(sizeof(double) * res_num * ntap);

    for(int i = 0;i < nfft;++i){
        extract_data(data,sdata,res_num * nfft,i,1,nfft);
        extract_data(coeffs,scoeff,ntap * nfft,i,1,nfft);

        lfilter(sdata,scoeff,out + i * res_num,res_num,ntap);
    }

    double *in = calloc(sizeof(double) , nfft * 2);
    fftw_complex *oout = malloc(sizeof(fftw_complex) * nfft);

    fftw_plan_with_nthreads(4);
    fftw_plan plan = fftw_plan_dft_1d(nfft,(fftw_complex *)in,oout,FFTW_BACKWARD,FFTW_ESTIMATE);
    if(plan == NULL) {
        printf("ifft:Create fftw_plan error!");
        exit(1);
    }
    
    for(size_t i = 0;i < res_num;++i){
        // printf("%ld\n",i);
        extract_data(out,in,res_num*nfft,i,2,res_num);
        // memcpy(in,temp,nfft * sizeof(double));
        fftw_execute(plan);
        memcpy(subband+i*nfft,oout,nfft * sizeof(fftw_complex));

    }

    // for(size_t i = 0;i < nfft;++i){
    //     printf("%f\n",oout[i][0]);
    // }
    

    fftw_destroy_plan(plan);
    fftw_cleanup_threads();

    free(sdata);
    free(scoeff);
    free(out);
}


void fft_convolve(double *d1,double *d2,double *out,size_t s1,size_t s2)
{
    size_t len = s1 + s2 - 1;
    double *t1 = malloc(sizeof(double) * len);
    double *t2 = malloc(sizeof(double) * len);
    double max = 0;
    fftw_complex *freq  = malloc(sizeof(fftw_complex) * len);
    fftw_complex *freq1 = malloc(sizeof(fftw_complex) * len);
    fftw_complex *freq2 = malloc(sizeof(fftw_complex) * len);

    for(size_t i = 0;i < len;++i){
        if(i < s1)
            t1[i] = d1[i];
        else
            t1[i] = 0;
    }

    for(size_t i = 0;i < len;++i){
        if(i < s2)
            t2[i] = d2[i];
        else
            t2[i] = 0;
    }

    fft_r2c(t1,freq1,len);
    fft_r2c(t2,freq2,len);
    for(size_t i = 0;i < len;++i){

        freq[i][0] = freq1[i][0] * freq2[i][0] - freq1[i][1] * freq2[i][1]; 
        freq[i][1] = freq1[i][0] * freq2[i][1] + freq1[i][1] * freq2[i][0];
    }

    fft_c2r(freq,out,len);

    for(int i = 0;i < len;++i){
        out[i] /= len;
    }

    free(t1);
    free(t2);
    fftw_free(freq);
    fftw_free(freq1);
    fftw_free(freq2);
}

void lfilter(double *data,double *coeffs,double *out,size_t l1,size_t l2)
{
    double *res = malloc(sizeof(double) * (l1+l2-1));
    fft_convolve(data,coeffs,res,l1,l2);
    memcpy(out,res,l1 * sizeof(double));
    free(res);
}